
pub mod base;
#[cfg(all(any(target_arch="x86", target_arch="x86_64"),feature="simd"))]
pub mod x86;
#[cfg(all(any(target_arch="arm", target_arch="aarch64"),feature="simd"))]
pub mod aarch64;

pub use base::*;
#[cfg(all(any(target_arch="x86", target_arch="x86_64"),feature="simd"))]
pub use x86::*;
#[cfg(all(any(target_arch="arm", target_arch="aarch64"),feature="simd"))]
pub use aarch64::*;

use crate::sha256::Sha256Ctx;

const MH_SHA256_H0:u32 = 0x6a09e667;
const MH_SHA256_H1:u32 = 0xbb67ae85;
const MH_SHA256_H2:u32 = 0x3c6ef372;
const MH_SHA256_H3:u32 = 0xa54ff53a;
const MH_SHA256_H4:u32 = 0x510e527f;
const MH_SHA256_H5:u32 = 0x9b05688c;
const MH_SHA256_H6:u32 = 0x1f83d9ab;
const MH_SHA256_H7:u32 = 0x5be0cd19;
 
// External Interface Definition
const ISAL_HASH_SEGS:usize = 16;
const ISAL_SHA256_BLOCK_SIZE:usize = 64;
const ISAL_MH_SHA256_BLOCK_SIZE:usize = ISAL_HASH_SEGS * ISAL_SHA256_BLOCK_SIZE;
const ISAL_SHA256_DIGEST_WORDS:usize = 8;

pub trait MhSha256BlockFunction {
    type FrameBuffer:Default;
    fn mh_sha256_block_function(input_data: &[u8], mh_sha256_segs_digests: &mut [[u32;ISAL_HASH_SEGS];ISAL_SHA256_DIGEST_WORDS], aligned_frame_buffer: &mut Self::FrameBuffer, num_blocks: u32);
}

#[repr(C)]
#[derive(Clone, Debug)]
pub struct IsalMhSha256Ctx<T:MhSha256BlockFunction> {
    _method_marker:std::marker::PhantomData<T>,
    // the digest of mti-hash SHA256
    mh_sha256_digest:[u32;ISAL_SHA256_DIGEST_WORDS], 
    //  Parameters for update feature, describe the lengths of input buffers in bytes
    total_length:usize,
    //  Padding the tail of input data for SHA256
    pub partial_block_buffer:[u8;ISAL_MH_SHA256_BLOCK_SIZE * 2],
    //  Storing the SHA256 interim digests of  all 16 segments. Each time, it will be copied to stack for 64-byte alignment purpose.
    mh_sha256_interim_digests:[[u32;ISAL_HASH_SEGS];ISAL_SHA256_DIGEST_WORDS],
    //  Re-structure sha256 block data from different segments to fit big endian. Use ISAL_AVX512_ALIGNED for 64-byte alignment purpose.
    frame_buffer:T::FrameBuffer,
}

impl<T:MhSha256BlockFunction> IsalMhSha256Ctx<T> {
    pub fn new() -> Self {
        Self {
            _method_marker:std::marker::PhantomData::<T>::default(),
            mh_sha256_digest:[0u32;ISAL_SHA256_DIGEST_WORDS], 
            total_length:0usize,
            partial_block_buffer:[0u8;ISAL_MH_SHA256_BLOCK_SIZE * 2],
            mh_sha256_interim_digests:seq!(N in 0..8 { [#([MH_SHA256_H~N;ISAL_HASH_SEGS],)*] }),
            frame_buffer:T::FrameBuffer::default(),
        }
    }

    pub fn update(&mut self, data:&[u8]) {
        let mh_sha256_segs_digests = &mut self.mh_sha256_interim_digests;
        let aligned_frame_buffer =  &mut self.frame_buffer;  // align 64
        let mut input_data = data;

        if input_data.len() == 0 { return; }
    
        let partial_block_len = self.total_length % ISAL_MH_SHA256_BLOCK_SIZE;
    
        self.total_length += input_data.len();
    
        // No enough input data for mh_sha256 calculation
        if input_data.len() + partial_block_len < ISAL_MH_SHA256_BLOCK_SIZE {
            self.partial_block_buffer[partial_block_len..partial_block_len + input_data.len()].copy_from_slice(input_data);
            return;
        }
    
        // mh_sha256 calculation for the previous partial block
        if partial_block_len != 0 {
            let copy_len = ISAL_MH_SHA256_BLOCK_SIZE - partial_block_len;
            self.partial_block_buffer[partial_block_len..partial_block_len + copy_len].copy_from_slice(&input_data[0..copy_len]);
            // Do one_block process
            T::mh_sha256_block_function(&self.partial_block_buffer, mh_sha256_segs_digests, aligned_frame_buffer, 1);
            input_data = &input_data[copy_len..];
        }
    
        // Calculate mh_sha256 for the current blocks
        let num_blocks = input_data.len() / ISAL_MH_SHA256_BLOCK_SIZE;
        if num_blocks > 0 {
            // Do num_blocks process
            T::mh_sha256_block_function(input_data, mh_sha256_segs_digests, aligned_frame_buffer, num_blocks as u32);
            input_data = &input_data[num_blocks * ISAL_MH_SHA256_BLOCK_SIZE..];
        }
    
        // Store the partial block
        if !input_data.is_empty() {
            self.partial_block_buffer[0..input_data.len()].copy_from_slice(input_data);
        }
    }

    pub fn finalize(&mut self, digest:&mut [u8;std::mem::size_of::<u32>() * ISAL_SHA256_DIGEST_WORDS]) {
        let mut partial_buffer_len = self.total_length % ISAL_MH_SHA256_BLOCK_SIZE;
        let frame_buffer = &mut self.frame_buffer;

        // Padding the first block
        self.partial_block_buffer[partial_buffer_len] = 0x80;
        partial_buffer_len += 1;

        self.partial_block_buffer[partial_buffer_len..ISAL_MH_SHA256_BLOCK_SIZE].fill(0);

        // Calculate the first block without total_length if padding needs 2 block
        if partial_buffer_len > (ISAL_MH_SHA256_BLOCK_SIZE - 8) {
            T::mh_sha256_block_function(&mut self.partial_block_buffer, &mut self.mh_sha256_interim_digests, frame_buffer, 1);
            // Padding the second block
            self.partial_block_buffer[0..ISAL_SHA256_BLOCK_SIZE].fill(0);
        }
        // Padding the block
        let len_in_bit:[u8;std::mem::size_of::<u64>()] = (self.total_length as u64 * 8).to_be_bytes();
        self.partial_block_buffer[ISAL_MH_SHA256_BLOCK_SIZE - 8..ISAL_MH_SHA256_BLOCK_SIZE].copy_from_slice(&len_in_bit);
        T::mh_sha256_block_function(&mut self.partial_block_buffer, &mut self.mh_sha256_interim_digests, frame_buffer, 1);

        let mh_sha256_interim_digests = unsafe { std::mem::transmute::<&[[u32;ISAL_HASH_SEGS];ISAL_SHA256_DIGEST_WORDS],&[u8;std::mem::size_of::<u32>()*ISAL_HASH_SEGS*ISAL_SHA256_DIGEST_WORDS]>(&self.mh_sha256_interim_digests) };
        // Calculate multi-hash SHA256 digests (segment digests as input message)
        Self::sha256_for_mh_sha256(mh_sha256_interim_digests, &mut self.mh_sha256_digest);

        /* Output the digests of mh_sha256 */
        seq!(N in 0..8 { digest[N*std::mem::size_of::<u32>()..(N+1)*std::mem::size_of::<u32>()].copy_from_slice(&self.mh_sha256_digest[N].to_le_bytes()); });
    }

    fn sha256_for_mh_sha256(input_data:&[u8], digest:&mut [u32;ISAL_SHA256_DIGEST_WORDS]) {
        let mut input_data = input_data;
        let origin_len = input_data.len();
        let mut buf = [0;2 * ISAL_SHA256_BLOCK_SIZE];

        digest[0] = MH_SHA256_H0;
        digest[1] = MH_SHA256_H1;
        digest[2] = MH_SHA256_H2;
        digest[3] = MH_SHA256_H3;
        digest[4] = MH_SHA256_H4;        
        digest[5] = MH_SHA256_H5;
        digest[6] = MH_SHA256_H6;
        digest[7] = MH_SHA256_H7;

        while input_data.len() >= ISAL_SHA256_BLOCK_SIZE {
                Sha256Ctx::sha256_single(input_data, digest);
                input_data = &input_data[ISAL_SHA256_BLOCK_SIZE..];
        }
        buf[0..input_data.len()].copy_from_slice(input_data);
        buf[input_data.len()] = 0x80;
        let buf_len = if input_data.len() + 1 > ISAL_SHA256_BLOCK_SIZE - 8 {
            2 * ISAL_SHA256_BLOCK_SIZE
        } else {
            ISAL_SHA256_BLOCK_SIZE
        };
        buf[input_data.len()+1..buf_len-8].fill(0);

        let data_bits_len = (origin_len as u64 * 8).to_be_bytes();
        buf[buf_len-8..buf_len].copy_from_slice(&data_bits_len);

        Sha256Ctx::sha256_single(&buf[0..ISAL_SHA256_BLOCK_SIZE], digest);
        if buf_len == 2 * ISAL_SHA256_BLOCK_SIZE {
            Sha256Ctx::sha256_single(&buf[ISAL_SHA256_BLOCK_SIZE..], digest);
        }
    }
}